/*
 * 内核DNS过滤模块，避免DNS被劫持
 * DNS响应会被随机替换，某些响应的IP被改为FakeIP
 * 1. 丢掉解析IP中包括FakeIP的响应
 *   insmod filtedns.ko address=1.2.3.4
 * 2. 有时FakeIP有正常的域名，如下参数可以丢掉包含FakeIP的响应，又避免丢掉它正常的响应
 *   insmod filtedns.ko address=1.2.3.4 domain=www.xxx.com
 */

#include "dns_header.h"

#include <linux/kernel.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <linux/tcp.h>
#include <linux/inet.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/version.h>
#include <linux/module.h>

#define DNS_MAXSIZE 255
#define DNS_LABEL_MAXSIZE 63
#define DNS_PORT 53

#define TYPE_A 0x0001
#define CLASS_IN 0x0001

#ifdef DEBUG
#define DEBUG_PRINT(fmt, ...)                                                  \
	{                                                                          \
		printk(KERN_DEBUG "xt_dns %s(%d):" fmt "\n", __func__, __LINE__,       \
		       ##__VA_ARGS__);                                                 \
	}
#else
#define DEBUG_PRINT(...)
#endif


static unsigned int  match_addr = 0;

static char match_ipaddr[16];
module_param_string(address, match_ipaddr, 16, 0);
MODULE_PARM_DESC(address, "The ip address of faked dns.");

static char match_domain[DNS_MAXSIZE];
module_param_string(domain, match_domain, DNS_MAXSIZE, 0);
MODULE_PARM_DESC(domain, "The real domain of address.");


int filter_header(struct sk_buff* skb)
{
	int offset = 0;
	const struct iphdr* iph;

	iph = ip_hdr(skb);
	if (iph->frag_off != 0)
		return 0;

	offset = iph->ihl * 4;

	switch (iph->protocol)
	{
	case IPPROTO_UDP:
	{
		const struct udphdr* uh;
		struct udphdr _udph;
		uh = skb_header_pointer(skb, offset, sizeof(_udph), &_udph);
		if (uh == NULL)
		{
			DEBUG_PRINT("filtedns: invalid udp header.");
			return 0;
		}
		if (uh->source != htons(DNS_PORT))
			return 0;

		return offset + sizeof(_udph);
	}
	case IPPROTO_TCP:
	{
		const struct tcphdr* th;
		struct tcphdr _tcph;
		th = skb_header_pointer(skb, offset, sizeof(_tcph), &_tcph);
		if (th == NULL)
		{
			DEBUG_PRINT("filtedns: invalid tcp header.");
			return 0;
		}

		if (!(th->ack & th->psh) || (th->source != htons(DNS_PORT)))
		{
			//DEBUG_PRINT("filtedns: not dns packet");
			return 0;
		}

		return offset + th->doff * 4 + 2;
	}
	default:
		break;
	}

	return 0;
}

int parse_dns_resp_domain(struct sk_buff* skb, int offset, uint8_t* qname)
{
	uint16_t qtype;                 // qtype buffer
	uint16_t klass;                 // class buffer

	uint16_t qlen; // qname length, MAX 255
	uint8_t llen;  // label length, MAX 63

	qlen = 0;
	llen = 255;
	while (llen != 0 && qlen < DNS_MAXSIZE)
	{
		// read label size
		if (skb_copy_bits(skb, offset, &llen, sizeof(uint8_t)) < 0 || llen > DNS_LABEL_MAXSIZE)
		{
			DEBUG_PRINT("filtedns: invalid label len %u->%x.", offset, llen);
			return 0;
		}

		offset += 1;
		if (llen == 0)
			break;

		if (qlen + llen + 1 <= DNS_MAXSIZE && skb_copy_bits(skb, offset, (qname + qlen), sizeof(uint8_t) * llen) < 0)
		{
			DEBUG_PRINT("filtedns: invalid label name %u,%u", qlen, llen);
			return 0;
		}
		offset += llen;
		qlen += llen;
		qname[qlen] = '.';
		qlen += 1;
	}

	if (qlen)
		qname[qlen - 1] = 0;
	DEBUG_PRINT("filtedns: success qname parse.[%s] %d", qname, qlen);

	if (skb_copy_bits(skb, offset, &qtype, sizeof(qtype)) < 0)
	{
		DEBUG_PRINT("filtedns: invalid qtype");
		return 0;
	}
	offset += 2;

	if (skb_copy_bits(skb, offset, &klass, sizeof(klass)) < 0)
	{
		DEBUG_PRINT("filtedns: invalid class");
		return 0;
	}
	offset += 2;

	if (ntohs(qtype) != TYPE_A || ntohs(klass) != CLASS_IN)
	{
		DEBUG_PRINT("filtedns: invalid type or class %x %x", qtype, klass);
		return 0;
	}

	return offset;
}

int parse_dns_resp_address(struct sk_buff* skb, int offset, uint32_t* addr)
{
	struct dns_resp_addr_hdr _ah;
	struct dns_resp_addr_hdr* ah;

	//DEBUG_PRINT("filtedns: parsing address. %u %d", offset, sizeof(_ah));
	ah = &_ah;
	if (skb_copy_bits(skb, offset, &_ah, sizeof(_ah)) < 0)
	{
		DEBUG_PRINT("filtedns: invalid answer %d %d", offset, sizeof(_ah));
		return 0;
	}

	//DEBUG_PRINT("filtedns: parsing address. %x %x %x %x %x", ah->name, ah->qtype, ah->klass, ah->ttl, ah->dlen);

	if (*(unsigned char*)ah != 0xc0)
	{
		DEBUG_PRINT("filtedns: invalid name %x", ah->name);
		return 0;
	}

	if (ntohs(ah->qtype) == TYPE_A && ntohs(ah->klass) == CLASS_IN)
	{
		if (skb_copy_bits(skb, offset + sizeof(_ah), addr, sizeof(*addr)) < 0)
		{
			DEBUG_PRINT("filtedns: invalid address");
			return 0;
		}
		DEBUG_PRINT("filtedns: address parsed. %u %x", *addr, ntohs(ah->dlen));
	}

	//DEBUG_PRINT("filtedns: address skip. %x", ntohs(ah->dlen));
	return offset + sizeof(_ah) + ntohs(ah->dlen);
}

// TODO: judge by kernel version
#if 0 // for lower kernel
unsigned int filter_dns_resp(unsigned int hooknum,
                             struct sk_buff* skb,
                             const struct net_device* in,
                             const struct net_device* out,
                             int (*okfn)(struct sk_buff*))

#else
unsigned int filter_dns_resp(void* priv,
                             struct sk_buff* skb,
                             const struct nf_hook_state* state)
#endif
{
	int offset;
	const struct dns_hdr* dh;
	struct dns_hdr _dnsh;

	int num_a;
	uint32_t addr;
	uint8_t qname[DNS_MAXSIZE]; // qname buffer

	offset = filter_header(skb);
	if (offset == 0)
		return NF_ACCEPT;

	dh = skb_header_pointer(skb, offset, sizeof(_dnsh), &_dnsh);

	// if dns query name != fiberhome && dns respone addr == 61.
	if (dh->qr != 1 || dh->opcode != 0 || dh->rcode != 0 || dh->ancount == 0)
		return NF_ACCEPT;

	offset += sizeof(_dnsh);
	offset = parse_dns_resp_domain(skb, offset, qname);
	if (offset == 0)
		return NF_ACCEPT;

	num_a = ntohs(dh->ancount);
	//DEBUG_PRINT("filtedns: number of answer: %d!", num_a);
	while (num_a -- > 0)
	{
		offset = parse_dns_resp_address(skb, offset, &addr);
		if (offset == 0)
			return NF_ACCEPT;

		// match
		if (match_addr == addr && ( match_domain[0] == 0 || strcmp(match_domain, (char*)qname) ))
		{
			DEBUG_PRINT("filtedns: MATCHEED!");
			return NF_DROP;
		}
	}

	return NF_ACCEPT;
}

static struct nf_hook_ops nfho =
{
	.hook = filter_dns_resp,
	.pf = PF_INET,
	.hooknum = NF_INET_LOCAL_IN,
	.priority = NF_IP_PRI_LAST
};

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4,13,0)
#define NF_REGISTER_NET_HOOK(reg) nf_register_net_hook(&init_net, reg)
#define NF_UNREGISTER_NET_HOOK(reg) nf_unregister_net_hook(&init_net, reg);
#else
#define NF_REGISTER_NET_HOOK(reg) nf_register_hook(reg)
#define NF_UNREGISTER_NET_HOOK(reg) nf_unregister_hook(reg);
#endif

static int __init sknf_init(void)
{
	if (NF_REGISTER_NET_HOOK(&nfho))
	{
		printk(KERN_ERR"nf_register_net_hook() failed\n");
        return -1;  
    }

    match_addr = in_aton(match_ipaddr);

    return 0;  
}  
  
static void __exit sknf_exit(void)  
{  
    NF_UNREGISTER_NET_HOOK(&nfho);  
}  
  
module_init(sknf_init);  
module_exit(sknf_exit);  
MODULE_LICENSE("GPL");

